import torch


def torch2onnx(
        weight: str = '',
        output_file: str = '',
        in_chans: int = 3,
        img_size: int = 480):

    DEVICE = torch.device("cpu")
    # 加载模型
    model = torch.load(weight)
    # 切换至eval模式
    model.eval().to(DEVICE)

    input_shape = (1, in_chans, img_size, img_size)
    class_img = torch.randn(input_shape).to(DEVICE)
    defect_img = torch.randn(input_shape).to(DEVICE)
    dummy_inputs = (class_img, defect_img)

    # 示例输入（假设输入为3通道224x224图像）
    example_input = torch.randn(1, in_chans, img_size, img_size)  # batch_size=1

    torch.onnx.export(
        model,  # 要导出的模型
        example_input,  # 示例输入
        output_file,  # 输出文件名
        input_names=["input"],  # 输入节点名称
        output_names=["output"],  # 输出节点名称
        dynamic_axes={  # 动态维度配置（如可变batch_size）
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        },
        opset_version=11  # ONNX操作集版本（推荐≥11）
    )



if __name__ == '__main__':
    weight = "convnext_s.pth"
    output_file = "convnext_s_224.onnx"
    torch2onnx(img_size=224, weight=weight, output_file=output_file)
